package com.agent.server.grpc;

import com.agent.common.protocol.AgentServiceGrpc;
import com.agent.common.protocol.AgentServiceOuterClass.AgentMessageRequest;
import com.agent.common.protocol.AgentServiceOuterClass.AgentMessageResponse;
import com.agent.common.protocol.AgentServiceOuterClass.AgentMetadata;
import com.agent.server.service.LlmService;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Service;

import java.util.UUID;
import java.util.concurrent.atomic.AtomicBoolean;

@Service
public class AgentServiceImpl extends AgentServiceGrpc.AgentServiceImplBase {

    private static final Logger logger = LoggerFactory.getLogger(AgentServiceImpl.class);
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    // We need to get the LlmService from the Spring context
    private static ApplicationContext applicationContext;
    
    @Autowired
    public void setApplicationContext(ApplicationContext context) {
        applicationContext = context;
    }
    
    private LlmService getLlmService() {
        return applicationContext.getBean(LlmService.class);
    }

    @Override
    public StreamObserver<AgentMessageRequest> streamMessages(StreamObserver<AgentMessageResponse> responseObserver) {
        logger.info("Received new gRPC streaming connection");
        
        // Create a unique session ID for this connection
        final String sessionId = UUID.randomUUID().toString();
        final AtomicBoolean completed = new AtomicBoolean(false);
        
        // Return an observer for client messages
        return new StreamObserver<>() {
            @Override
            public void onNext(AgentMessageRequest request) {
                logger.info("Received client request: {}", request.getPayload());
                
                // Process the client message by sending to LLM and streaming responses
                if (getLlmService().streamFromLlm(sessionId, request.getPayload(), sseEvent -> {
                    try {
                        // Parse the SSE event
                        if (sseEvent.startsWith("data:")) {
                            // Extract the JSON data after "data:"
                            String jsonData = sseEvent.substring(5).trim();
                            if (jsonData.equals("[DONE]")) {
                                // This is the end of the stream
                                if (completed.compareAndSet(false, true)) {
                                    getLlmService().closeConnection(sessionId);
                                    responseObserver.onCompleted();
                                }
                                return;
                            }
                            
                            JsonNode responseNode = objectMapper.readTree(jsonData);
                            
                            // Check for different SSE event formats
                            String content = extractContent(responseNode);
                            
                            // Check if this is the final message
                            boolean isComplete = isCompletionMessage(responseNode);
                            
                            // Build response metadata
                            AgentMetadata metadata = AgentMetadata.newBuilder()
                                    .setAgentId(request.getMetadata().getAgentId())
                                    .setType("llm_response")
                                    .setMessageId(UUID.randomUUID().toString())
                                    .setTimestamp(System.currentTimeMillis())
                                    .build();
                            
                            // Build and send response
                            AgentMessageResponse response = AgentMessageResponse.newBuilder()
                                    .setMetadata(metadata)
                                    .setPayload(content)
                                    .setStatus(isComplete ? 1 : 0)
                                    .build();
                            
                            responseObserver.onNext(response);
                            
                            // If this is the last message, mark completion
                            if (isComplete) {
                                if (completed.compareAndSet(false, true)) {
                                    getLlmService().closeConnection(sessionId);
                                    responseObserver.onCompleted();
                                }
                            }
                        }
                    } catch (JsonProcessingException e) {
                        logger.error("Error processing LLM response", e);
                    }
                })) {
                    logger.info("Started LLM streaming for session: {}", sessionId);
                } else {
                    logger.error("Failed to start LLM streaming for session: {}", sessionId);
                    if (completed.compareAndSet(false, true)) {
                        sendErrorResponse(responseObserver, request, "Failed to connect to LLM service");
                        responseObserver.onCompleted();
                    }
                }
            }

            @Override
            public void onError(Throwable t) {
                logger.error("Error in client stream", t);
                if (completed.compareAndSet(false, true)) {
                    getLlmService().closeConnection(sessionId);
                    responseObserver.onCompleted();
                }
            }

            @Override
            public void onCompleted() {
                logger.info("Client completed stream");
                if (completed.compareAndSet(false, true)) {
                    getLlmService().closeConnection(sessionId);
                    responseObserver.onCompleted();
                }
            }
        };
    }
    
    /**
     * Extract content from the SSE event based on different possible formats
     */
    private String extractContent(JsonNode responseNode) {
        // Check common SSE response formats
        if (responseNode.has("content")) {
            return responseNode.get("content").asText("");
        } else if (responseNode.has("message") && responseNode.get("message").has("content")) {
            return responseNode.get("message").get("content").asText("");
        } else if (responseNode.has("choices") && responseNode.get("choices").size() > 0) {
            // OpenAI format
            JsonNode choice = responseNode.get("choices").get(0);
            if (choice.has("delta") && choice.get("delta").has("content")) {
                return choice.get("delta").get("content").asText("");
            } else if (choice.has("text")) {
                return choice.get("text").asText("");
            }
        }
        
        // Default to empty string if no content is found
        return "";
    }
    
    /**
     * Determine if this is a completion message (last message in the stream)
     */
    private boolean isCompletionMessage(JsonNode responseNode) {
        if (responseNode.has("is_complete")) {
            return responseNode.get("is_complete").asBoolean();
        } else if (responseNode.has("choices") && responseNode.get("choices").size() > 0) {
            // OpenAI format
            JsonNode choice = responseNode.get("choices").get(0);
            if (choice.has("finish_reason") && !choice.get("finish_reason").isNull()) {
                return true;
            }
        }
        
        return false;
    }
    
    private void sendErrorResponse(StreamObserver<AgentMessageResponse> responseObserver, 
                                  AgentMessageRequest request,
                                  String errorMessage) {
        AgentMetadata metadata = AgentMetadata.newBuilder()
                .setAgentId(request.getMetadata().getAgentId())
                .setType("error")
                .setMessageId(UUID.randomUUID().toString())
                .setTimestamp(System.currentTimeMillis())
                .build();
        
        AgentMessageResponse response = AgentMessageResponse.newBuilder()
                .setMetadata(metadata)
                .setPayload(errorMessage)
                .setStatus(-1)
                .build();
        
        responseObserver.onNext(response);
    }
} 